"""
@author:王耀
@date:2021/9/22
"""
import numpy as np
from queue import Queue
import math
import cv2
import otsu

def isNeighborConnect(nmsmatrix, w, h,th1, th2):
    """:cvar
    作用：判断是否和周围像素连接
    nmsmatrix: 非极大值抑制后的矩阵
    w， h：图片的宽和高
    th1，th2：阈值

    return: 1：如果周围存在大于th2；0：其他
    """
    if nmsmatrix[w - 1, h] > th2 or nmsmatrix[w + 1, h] > th2 or nmsmatrix[w, h - 1] > th2 or nmsmatrix[w, h + 1] > th2 \
        or nmsmatrix[w + 1, h + 1] > th2 or nmsmatrix[w + 1, h - 1] > th2 or nmsmatrix[w - 1, h + 1] > th2 or \
            nmsmatrix[w - 1, h - 1] > th2:
        return 1

    return 0



def nms(grad, direct):
    """:cvar
    作用：传统方法非极大值抑制
    grad：梯度矩阵
    direct：方向矩阵

    返回：非极大值抑制后的矩阵
    """
    # 非极大值抑制部分
    W, H = grad.shape
    nmsmatrix = np.copy(grad)
    nmsmatrix[0, :] = nmsmatrix[W - 1, :] = nmsmatrix[:, 0] = nmsmatrix[:, H - 1] = 0
    """
    如果当前像素本身就是0，则nmmsmatrix直接为0。
    否则判断梯度方向，将梯度方向分成四个方向，并对比方向上领域与当前像素关系。
    只有当前像素大于两边的像素才保留，否则置零。
    """
    for w in range(1, W - 1):
        for h in range(1, H - 1):
            if grad[w, h] == 0:
                nmsmatrix[w, h] = 0
            else:
                if (direct[w, h] <= 22.5 and direct[w, h] >= 0) or (direct[w, h] > 337.5):
                    neibor1 = grad[w, h - 1]
                    neibor2 = grad[w, h + 1]
                elif (direct[w, h] <= 67.5 and direct[w, h] > 22.5) or (direct[w, h] <= 247.5 and direct[w, h] > 202.5):
                    neibor1 = grad[w - 1, h - 1]
                    neibor2 = grad[w + 1, h + 1]
                elif (direct[w, h] <= 112.5 and direct[w, h] > 67.5) or (direct[w, h] <= 292.5 and direct[w, h] > 247.5):
                    neibor1 = grad[w - 1, h]
                    neibor2 = grad[w + 1, h]
                elif (direct[w, h] <= 157.5 and direct[w, h] > 112.5) or (direct[w, h] <= 337.5 and direct[w, h] > 292.5):
                    neibor1 = grad[w + 1, h - 1]
                    neibor2 = grad[w - 1, h + 1]
                else:
                    neibor1 = grad[w, h - 1]
                    neibor2 = grad[w, h + 1]

                if nmsmatrix[w, h] < neibor1 or nmsmatrix[w, h] < neibor2:
                    nmsmatrix[w, h] = 0

    # cv2.imshow('nms', np.copy(nmsmatrix).astype(np.uint8))
    # cv2.imwrite('nms.png', np.copy(nmsmatrix).astype(np.uint8))
    # cv2.waitKey()

    print('执行完毕：传统NMS算法')

    return nmsmatrix



def nmsWithLinearInter(grad, gradx, grady, direct):
    """:cvar
    作用：插值处理方法非极大值抑制
    grad：梯度矩阵
    gradx：x轴方向梯度矩阵
    grady：y轴方向梯度矩阵
    direct：方向矩阵

    返回：非极大值抑制后的矩阵
    """
    # 非极大值抑制部分
    W, H = grad.shape
    nmsmatrix = np.copy(grad)
    nmsmatrix[0, :] = nmsmatrix[W - 1, :] = nmsmatrix[:, 0] = nmsmatrix[:, H - 1] = 0
    """
    插值操作。
    如果gradx>grady，代表在x轴上变化幅度大，grady/gradx作为权值；
    如果gradx<grady，代表在y轴上变化幅度大，gradx/grady作为权值；
    同时要比梯度在xy轴上的方向，找到应当用哪几个领域作插值计算
    """
    for w in range(1, W - 1):
        for h in range(1, H - 1):
            neibor1 = weight = neibor2 = neibor3 = neibor4 = tmp = None
            if grad[w, h] == 0:
                nmsmatrix[w, h] = 0
            else:
                tmpx = gradx[w, h]
                tmpy = grady[w, h]
                tmp = grad[w, h]
                # 如果在y方向上的幅度比较大
                if np.abs(tmpy) > np.abs(tmpx):
                    weight = 1 / direct[w, h]  # 此时的插值的权值选取dx/dy，因此是1/direct[w,h]
                    neibor2 = grad[w - 1, h]
                    neibor4 = grad[w + 1, h]
                    if tmpx * tmpy > 0:  # 梯度在x、y轴方向符号相同
                        neibor1 = grad[w - 1, h - 1]
                        neibor3 = grad[w + 1, h + 1]
                    else:  # 符号相反
                        neibor1 = grad[w - 1, h + 1]
                        neibor3 = grad[w + 1, h - 1]
                # 在x轴方向幅度较大
                else:
                    weight = direct[w, h]
                    neibor2 = grad[w, h - 1]
                    neibor4 = grad[w, h + 1]
                    if tmpx * tmpy > 0:  # 梯度在x、y轴方向符号相同
                        neibor1 = grad[w + 1, h - 1]
                        neibor3 = grad[w - 1, h + 1]
                    else:  # 符号相反
                        neibor1 = grad[w - 1, h - 1]
                        neibor3 = grad[w - 1, h + 1]

                # 计算两个插值
                resultgrad1 = weight * neibor1 + (1 - weight) * neibor2
                resultgrad2 = weight * neibor3 + (1 - weight) * neibor4
                if tmp >= resultgrad1 and tmp >= resultgrad2:
                    nmsmatrix[w, h] = tmp
                else:
                    nmsmatrix[w, h] = 0

    # cv2.imshow('nms_linear', np.copy(nmsmatrix).astype(np.uint8))
    # cv2.imwrite('nms_linear.png', np.copy(nmsmatrix).astype(np.uint8))
    # cv2.waitKey()

    print('执行完毕：插值处理NMS算法')

    return nmsmatrix



def nmsAndDualThreds(grad, gradx, grady, direct, FLAG = 0, thres1 = None, thres2 = None, isotsu = None):
    """
    作用：非极大值抑制+双阈值检测
    grad：梯度矩阵
    gradx:x方向上的梯度矩阵
    grady：y方向上的梯度矩阵
    direct：方向矩阵
    FLAG：默认为0，表示是否使用线性插值进行NMS
    thres1、thres2:双阈值检测阈值，默认None，如果没有参数传入，则调用otsu.otsuThred来获得自适应的阈值

    返回：Canny算法最后的图片。
    """
    # 是否用插值处理非极大值抑制
    if not FLAG:  # 传统方法的非极大值抑制
        nmsmatrix = nms(grad, direct)
    else:  # 考虑插值
        nmsmatrix = nmsWithLinearInter(grad, gradx, grady, direct)

    # 两个阈值必须都有数字才能有效，否则自动获取阈值
    if thres1 != None and thres2 != None:
        th1 = min(thres1, thres2)
        th2 = max(thres1, thres2)
    # 自动获取阈值有两种选择，如果isotsu==1，则用Otsu算法获得阈值，如果otsu==0，则用最大像素的0.1和0.3作为阈值
    else:
        if isotsu:
            print('Use otsu algorithm to get threds...')
            th1, th2 = otsu.otsuThred(grad)
        else:
            th1 = 0.1 * np.max(nmsmatrix)
            th2 = 0.3 * np.max(nmsmatrix)

    W, H = grad.shape

    """
    双阈值检测、边缘连接部分
    如果大于大阈值，则一定保留，小于小阈值，则一定抛弃
    如果灰度介于两阈值之间，则判断8领域是否有已经保留的阈值：
    如果有则保留该点，否则抛弃
    """
    resimg = np.zeros((W, H), dtype=np.uint8)
    for w in range(1, W - 1):
        for h in range(1, H - 1):
            if (nmsmatrix[w, h] < th1):
                resimg[w, h] = 0
            elif (nmsmatrix[w, h] > th2):
                resimg[w, h] = 255
            else:  # 边缘连接
                if isNeighborConnect(nmsmatrix, w, h, th1, th2) == 1:
                    resimg[w, h] = 255
                else:
                    resimg[w, h] = 0

    print('执行完毕：双阈值检测+边缘链接')

    return resimg